class Agent(nn.Module):
    def __init__(self, input_dim, output_dim):
        """
        每个智能体专门处理一个任务的特征并输出标签。
        Args:
        - input_dim: 输入特征的维度
        - output_dim: 输出标签维度
        """
        super(Agent, self).__init__()
        self.fc = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        return self.fc(x)

class MultiAgentSeq2Label(nn.Module):
    def __init__(self, input_dim, output_dims):
        """
        Args:
        - input_dim: 输入特征的维度（多模态特征）
        - output_dims: 每个任务的输出维度，列表形式，例如 [3, 3, 3] 表示 3 个任务，每个任务 3 个输出
        """
        super(MultiAgentSeq2Label, self).__init__()
        self.agents = nn.ModuleList([Agent(input_dim, output_dim) for output_dim in output_dims])

    def forward(self, x):
        """
        Args:
        - x: 输入特征 (batch_size, seq_length, input_dim)
        
        Returns:
        - 各个任务的输出，每个任务的输出形状为 (batch_size, seq_length, output_dim)
        """
        outputs = []
        for agent in self.agents:
            outputs.append(agent(x))  # 每个任务独立输出
        return outputs  # 返回每个任务的输出列表

# 示例：
input_dim = 768
output_dims = [3, 3, 3]  # 假设有三个任务，每个任务有3个输出
multi_agent_model = MultiAgentSeq2Label(input_dim, output_dims)

# 输入特征 (batch_size=32, seq_length=100, input_dim=768)
input_tensor = torch.randn(32, 100, 768)

# 前向传播，每个任务有独立的输出
task_outputs = multi_agent_model(input_tensor)
for i, task_output in enumerate(task_outputs):
    print(f"Task {i} output shape: {task_output.shape}")